import baostock as bs
import pandas as pd
import baostock.common.context as conx


def _query_basic_stock(code):
    rs = bs.query_stock_basic(code=code)
    result = {}
    if (rs.error_code == '0') & rs.next():
        data = rs.get_row_data()
        for i in range(0, len(data)):
            result[rs.fields[i]] = data[i]
    return result

def _history_k_data(code, start_date, end_date,
                    fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,"
                           "isST", frequency="d"):
    rs = bs.query_history_k_data_plus(code,
                                      fields,
                                      start_date=start_date, end_date=end_date,
                                      frequency=frequency, adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    data = pd.DataFrame(data_list, columns=rs.fields)
    return data


def _query_stock_industry(code="", date=""):
    rs = bs.query_stock_industry(code, date)
    data_list = []
    while (rs.error_code == '0') & rs.next():
        data_list.append(rs.get_row_data())
    data = pd.DataFrame(data_list, columns=rs.fields)
    return data


def _query_industry_one_by_one(code="", date=""):
    rs = bs.query_stock_industry(code, date)
    while (rs.error_code == '0') & rs.next():
        yield rs.get_row_data()


def _history_k_data_one_by_one(code, start_date, end_date,
                    fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,"
                           "isST", frequency="d"):
    rs = bs.query_history_k_data_plus(code,
                                      fields,
                                      start_date=start_date, end_date=end_date,
                                      frequency=frequency, adjustflag="3")
    while (rs.error_code == '0') & rs.next():
        yield rs.get_row_data()


class QueryStock(object):
    def __init__(self, user_id='anonymous', password='123456', options=0):
        self.user_id = user_id
        self.password = password
        self.options = options
        self.default_socket = None
        self.login = False

    def __enter__(self):
        data = bs.login(self.user_id, self.password, self.options)
        if data.error_code != '0':
            raise Exception('登陆失败，with message %s' % data.error_msg)
        if hasattr(conx, 'default_socket'):
            print('登录成功')
            self.default_socket = getattr(conx, 'default_socket')
            self.login = True
        return self

    def __login__(self):
        self.__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        if hasattr(conx, 'default_socket'):
            bs.logout()
            self.default_socket = None
            self.login = False
        return False

    def logout(self):
        if hasattr(conx, 'default_socket'):
            bs.logout()
            self.default_socket = None
            self.login = False
        return False

    def _handle_query_data(self, _call=None, args=(), kwargs=None):
        if kwargs is None:
            kwargs = {}
        if self.login is False:
            raise Exception('请先登录')
        elif self.default_socket is None:
            raise Exception('未知错误')
        elif not hasattr(conx, 'default_socket'):
            setattr(conx, 'default_socket', self.default_socket)
        if hasattr(_call, '__call__'):
            return _call(*args, **kwargs)
        else:
            print('handler 必须是可执行的')
            return None

    def check(self):
        if self.login is False:
            print('尚未登录')
        elif self.default_socket is not None:
            print('检测成功')
        else:
            print('检测失败')

    def query_k_data(self, code, start_date, end_date,
                     fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,"
                            "isST", frequency="d"
                     ):
        return self._handle_query_data(_call=_history_k_data, args=(code, start_date, end_date,fields,), kwargs={"frequency": frequency})

    def query_industry(self, code="", date=""):
        return self._handle_query_data(_call=_query_stock_industry, args=(code, date))

    def get_stock_info(self, code=""):
        return self._handle_query_data(_call=_query_basic_stock, args=(code, ))

    def query_industry_one_by_one(self, code="", date=""):
        return self._handle_query_data(_call=_query_industry_one_by_one, args=(code, date))

    def query_history_one_by_one(self,code, start_date, end_date,
                     fields="date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,"
                            "isST", frequency="d"):
        return self._handle_query_data(_call=_history_k_data_one_by_one, args=(code, start_date, end_date,fields,), kwargs={"frequency": frequency})


if __name__ == '__main__':
    with QueryStock() as q:
        print(q.get_stock_info(code='sh.600001'))
